# ******************************************************************************
#  Copyright (c) 2021-2022. Kneron Inc. All rights reserved.                   *
# ******************************************************************************

from .KPWrapperUtils import KPWrapperUtils as wrapper_utils
from .KPValue import DeviceDescriptorList, DeviceGroup, ModelNefDescriptor, SystemInfo, DdrManageAttributes
from .KPWrapper import KPWrapper
from .KPEnum import *
from .KPException import _check_api_return_code
from typing import List
import time
import ctypes


class KPCore:
    @staticmethod
    def scan_devices() -> DeviceDescriptorList:
        """
        Scan all Kneron devices and report a list.
        This function can get devices connectivity information at runtime.

        Returns
        -------
        device_descriptor_list : kp.DeviceDescriptorList
            DeviceDescriptorList object, contain information of connected devices from USB perspectives.

        See Also
        --------
        kp.DeviceDescriptorList
        """
        device_descriptor_list = DeviceDescriptorList()

        device_descriptor_list._element_buffer = KPWrapper().LIB.kp_scan_devices()
        device_descriptor_list._is_allocate_from_c = True

        return device_descriptor_list

    @staticmethod
    def connect_devices(usb_port_ids: List[int]) -> DeviceGroup:
        """
        To connect multiple (including one) Kneron devices.

        Parameters
        ----------
        usb_port_ids : List[int]
            An list contains device's port ID which can be known from scan_devices(), if [0] is given then it will
            try to connect first connectable device.

        Returns
        -------
        device_group : kp.DeviceGroup
            Represents a set of devices handle.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.scan_devices : Scan all Kneron devices and report a list.
        """
        num_devices = len(usb_port_ids)
        usb_port_ids = [int(usb_port_id) for usb_port_id in usb_port_ids]
        device_port_ids_c_array = (ctypes.c_int * len(usb_port_ids))(*usb_port_ids)
        device_port_ids_c_array = ctypes.cast(device_port_ids_c_array, ctypes.POINTER(ctypes.c_int))
        status = ctypes.c_int()

        device_group_pointer = KPWrapper().LIB.kp_connect_devices(num_devices,
                                                                  device_port_ids_c_array,
                                                                  status)

        device_group = DeviceGroup(address=device_group_pointer)
        device_group._is_allocate_from_c = True

        _check_api_return_code(result=device_group,
                               api_return_code=ApiReturnCode(status.value))

        return device_group

    @staticmethod
    def connect_devices_without_check(usb_port_ids: List[int]) -> DeviceGroup:
        """
        To connect multiple (including one) Kneron devices without any examinations of system info.

        Parameters
        ----------
        usb_port_ids : List[int]
            An list contains device's port ID which can be known from scan_devices(), if [0] is given then it will
            try to connect first connectable device.

        Returns
        -------
        device_group : kp.DeviceGroup
            Represents a set of devices handle.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.scan_devices : Scan all Kneron devices and report a list.
        """
        num_devices = len(usb_port_ids)
        usb_port_ids = [int(usb_port_id) for usb_port_id in usb_port_ids]
        device_port_ids_c_array = (ctypes.c_int * len(usb_port_ids))(*usb_port_ids)
        device_port_ids_c_array = ctypes.cast(device_port_ids_c_array, ctypes.POINTER(ctypes.c_int))
        status = ctypes.c_int()

        device_group_pointer = KPWrapper().LIB.kp_connect_devices_without_check(num_devices,
                                                                                device_port_ids_c_array,
                                                                                status)

        device_group = DeviceGroup(address=device_group_pointer)
        device_group._is_allocate_from_c = True

        _check_api_return_code(result=device_group,
                               api_return_code=ApiReturnCode(status.value))

        return device_group

    @staticmethod
    def disconnect_devices(device_group: DeviceGroup) -> None:
        """
        To disconnect a Kneron device.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        """
        status = ApiReturnCode.KP_SUCCESS.value

        if device_group._is_allocate_from_c:
            status = wrapper_utils.disconnect_devices_by_address(device_group_address=device_group.address)
            device_group._is_allocate_from_c = not (ApiReturnCode.KP_SUCCESS == ApiReturnCode(status))

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def set_timeout(device_group: DeviceGroup, milliseconds: int) -> None:
        """
        To set a global timeout value for all USB communications with the device.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        milliseconds : int
            Pre-set timeout value in milliseconds.

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        """
        KPWrapper().LIB.kp_set_timeout(device_group.address, milliseconds)

    @staticmethod
    def reset_device(device_group: DeviceGroup, reset_mode: ResetMode, sleep_secs: float = 0) -> None:
        """
        Reset the device in hardware mode or software mode.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        reset_mode : kp.ResetMode
            Refer to ResetMode.
        sleep_secs : float
            Set sleep time in seconds for reboot device.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.ResetMode
        """
        status = KPWrapper().LIB.kp_reset_device(device_group.address, reset_mode.value)
        time.sleep(sleep_secs)

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def load_firmware(device_group: DeviceGroup,
                      scpu_fw_buffer: bytes,
                      ncpu_fw_buffer: bytes) -> None:
        """
        Upload firmware data bytes to device through USB.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        scpu_fw_buffer : bytes
            A bytes buffer contains the content of SCPU firmware file.
        ncpu_fw_buffer : bytes
            A bytes buffer contains the content of NCPU firmware file.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_firmware_from_file : Upload firmware to device through USB by firmware file path.
        """
        scpu_fw_buffer_c_array = (ctypes.c_ubyte * len(scpu_fw_buffer)).from_buffer(bytearray(scpu_fw_buffer))
        ncpu_fw_buffer_c_array = (ctypes.c_ubyte * len(ncpu_fw_buffer)).from_buffer(bytearray(ncpu_fw_buffer))

        status = KPWrapper().LIB.kp_load_firmware(device_group.address,
                                                  ctypes.byref(scpu_fw_buffer_c_array),
                                                  len(scpu_fw_buffer),
                                                  ctypes.byref(ncpu_fw_buffer_c_array),
                                                  len(ncpu_fw_buffer))

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def load_firmware_from_file(device_group: DeviceGroup, scpu_fw_path: str, ncpu_fw_path: str) -> None:
        """
        Upload firmware to device through USB by firmware file path.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        scpu_fw_path : str
            SCPU firmware file path.
        ncpu_fw_path : str
            NCPU firmware file path.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_firmware : Upload firmware data bytes to device through USB.
        """
        status = KPWrapper().LIB.kp_load_firmware_from_file(device_group.address,
                                                            scpu_fw_path.encode('utf-8'),
                                                            ncpu_fw_path.encode('utf-8'))

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def load_model(device_group: DeviceGroup, nef_buffer: bytes) -> ModelNefDescriptor:
        """
        Upload NEF models data bytes to device through USB.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        nef_buffer : bytes
            A bytes buffer contains the content of NEF file.

        Returns
        -------
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_model_from_file : Upload NEF Model to device through USB by NEF file path.
        kp.core.get_model_info : Get model information (crc, model id, etc.).
        kp.ModelNefDescriptor
        """
        model_nef_descriptor = ModelNefDescriptor()
        model_buffer_c_array = (ctypes.c_ubyte * len(nef_buffer)).from_buffer(bytearray(nef_buffer))
        status = KPWrapper().LIB.kp_load_model(device_group.address,
                                               model_buffer_c_array,
                                               len(nef_buffer),
                                               ctypes.byref(model_nef_descriptor._get_element_buffer()))
        model_nef_descriptor._is_allocate_from_c = True

        _check_api_return_code(result=model_nef_descriptor,
                               api_return_code=ApiReturnCode(status))

        return model_nef_descriptor

    @staticmethod
    def load_model_from_file(device_group: DeviceGroup, file_path: str) -> ModelNefDescriptor:
        """
        Upload NEF Model to device through USB by NEF file path.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        file_path : str
            NEF model file path.

        Returns
        -------
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Raises
        ------
        ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_model : Upload NEF models data bytes to device through USB.
        kp.core.get_model_info : Get model information (crc, model id, etc.).
        kp.ModelNefDescriptor
        """
        model_nef_descriptor = ModelNefDescriptor()
        status = KPWrapper().LIB.kp_load_model_from_file(device_group.address,
                                                         file_path.encode('utf-8'),
                                                         ctypes.byref(model_nef_descriptor._get_element_buffer()))
        model_nef_descriptor._is_allocate_from_c = True

        _check_api_return_code(result=model_nef_descriptor,
                               api_return_code=ApiReturnCode(status))

        return model_nef_descriptor

    @staticmethod
    def load_encrypted_models(device_group: DeviceGroup, nef_buffer_list: List[bytes]) -> ModelNefDescriptor:
        """
        Upload encrypted NEF models data bytes to device through USB.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        nef_buffer_list : List[bytes]
            A list of byte buffer that contains the content of encrypted NEF file(s).

        Returns
        -------
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_encryted_models_from_file : Upload encryted NEF Model to device through USB by encrypted NEF file path(s).
        kp.core.get_model_info : Get model information (crc, model id, etc.).
        kp.ModelNefDescriptor
        """
        model_buffers_to_c_array = (ctypes.c_char_p * len(nef_buffer_list))(*nef_buffer_list)
        model_nef_descriptor = ModelNefDescriptor()
        status = KPWrapper().LIB.kp_load_encrypted_models(device_group.address,
                                                          model_buffers_to_c_array,
                                                          len(nef_buffer_list[0]),
                                                          len(nef_buffer_list),
                                                          ctypes.byref(model_nef_descriptor._get_element_buffer()))
        model_nef_descriptor._is_allocate_from_c = True

        _check_api_return_code(result=model_nef_descriptor,
                               api_return_code=ApiReturnCode(status))

        return model_nef_descriptor

    @staticmethod
    def load_encrypted_models_from_file(device_group: DeviceGroup,
                                        file_path_list: List[str]) -> ModelNefDescriptor:
        """
        Upload encrypted NEF Models to device through USB by encrypted NEF file path(s).

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        file_path_list : List[str]
            A list of encrypted NEF model file path(s).

        Returns
        -------
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Raises
        ------
        ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_encrypted_model : Upload encrypted NEF models data bytes to device through USB.
        kp.core.get_model_info : Get model information (crc, model id, etc.).
        kp.ModelNefDescriptor
        """
        file_path_list = [file_path.encode('utf-8') for file_path in file_path_list]
        file_paths_to_c = (ctypes.c_char_p * len(file_path_list))(*file_path_list)
        model_nef_descriptor = ModelNefDescriptor()
        status = KPWrapper().LIB.kp_load_encrypted_models_from_file(device_group.address,
                                                                    file_paths_to_c,
                                                                    len(file_path_list),
                                                                    ctypes.byref(
                                                                        model_nef_descriptor._get_element_buffer()))
        model_nef_descriptor._is_allocate_from_c = True

        _check_api_return_code(result=model_nef_descriptor,
                               api_return_code=ApiReturnCode(status))

        return model_nef_descriptor

    @staticmethod
    def enable_firmware_log(device_group: DeviceGroup, usb_port_id: int, log_file_path: str) -> None:
        """
        Enable firmware log from certain device.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        usb_port_id : int
            The device port ID to enable firmware log.
        log_file_path : str
            The log file output path.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.disable_firmware_log : Disable firmware log of all devices with firmware log enabled.
        """
        status = KPWrapper().LIB.kp_enable_firmware_log(device_group.address,
                                                        usb_port_id,
                                                        log_file_path.encode('utf-8'))

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def disable_firmware_log(device_group: DeviceGroup) -> None:
        """
        Disable firmware log of all devices with firmware log enabled.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.enable_firmware_log : Enable firmware log of all devices with firmware log enabled.
        """
        status = KPWrapper().LIB.kp_disable_firmware_log(device_group.address)

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def get_system_info(device_group: DeviceGroup, usb_port_id: int) -> SystemInfo:
        """
        Get system information (kn number and firmware version).

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        usb_port_id : int
            Specific device port ID.

        Returns
        -------
        system_info : kp.SystemInfo
            SystemInfo object for describing the system information of specific device.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.SystemInfo
        """
        system_info = SystemInfo()
        status = KPWrapper().LIB.kp_get_system_info(device_group.address,
                                                    usb_port_id,
                                                    ctypes.byref(system_info._get_element_buffer()))

        _check_api_return_code(result=system_info,
                               api_return_code=ApiReturnCode(status))

        return system_info

    @staticmethod
    def get_model_info(device_group: DeviceGroup, usb_port_id: int) -> ModelNefDescriptor:
        """
        Get model information (crc, model id, etc.).

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        usb_port_id : int
            Specific device port ID.

        Returns
        -------
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core.load_model : Upload NEF models data bytes to device through USB.
        kp.core.load_model_from_file : Upload NEF Model to device through USB by NEF file path.
        kp.ModelNefDescriptor
        """
        model_nef_descriptor = ModelNefDescriptor()
        status = KPWrapper().LIB.kp_get_model_info(device_group.address,
                                                   usb_port_id,
                                                   ctypes.byref(model_nef_descriptor._get_element_buffer()))
        model_nef_descriptor._is_allocate_from_c = True

        _check_api_return_code(result=model_nef_descriptor,
                               api_return_code=ApiReturnCode(status))

        return model_nef_descriptor

    @staticmethod
    def _update_kdp2_firmware(device_group: DeviceGroup,
                              scpu_fw_buffer: bytes,
                              ncpu_fw_buffer: bytes,
                              auto_reboot: bool) -> None:
        """
        (Advance) Update kdp2 firmware data bytes to device flash through USB.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        scpu_fw_buffer : bytes
            A bytes buffer contains the content of SCPU firmware file.
        ncpu_fw_buffer : bytes
            A bytes buffer contains the content of NCPU firmware file.
        auto_reboot : bool
            Auto reboot or not after update firmware.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core._update_kdp2_firmware_from_file : (Advance) Update kdp2 firmware to device flash through USB by firmware file path.
        """
        scpu_fw_buffer_c_array = (ctypes.c_ubyte * len(scpu_fw_buffer)).from_buffer(bytearray(scpu_fw_buffer))
        ncpu_fw_buffer_c_array = (ctypes.c_ubyte * len(ncpu_fw_buffer)).from_buffer(bytearray(ncpu_fw_buffer))

        status = KPWrapper().LIB.kp_update_kdp2_firmware(device_group.address,
                                                         ctypes.byref(scpu_fw_buffer_c_array),
                                                         len(scpu_fw_buffer),
                                                         ctypes.byref(ncpu_fw_buffer_c_array),
                                                         len(ncpu_fw_buffer),
                                                         auto_reboot)

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def _update_kdp2_firmware_from_file(device_group: DeviceGroup,
                                        scpu_fw_path: str,
                                        ncpu_fw_path: str,
                                        auto_reboot: bool) -> None:
        """
        (Advance) Update kdp2 firmware to device flash through USB by firmware file path.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        scpu_fw_path : str
            SCPU firmware file path.
        ncpu_fw_path : str
            NCPU firmware file path.
        auto_reboot : bool
            Auto reboot or not after update firmware.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.core._update_kdp2_firmware : (Advance) Update kdp2 firmware data bytes to device flash through USB.
        """
        status = KPWrapper().LIB.kp_update_kdp2_firmware_from_files(device_group.address,
                                                                    scpu_fw_path.encode('utf-8'),
                                                                    ncpu_fw_path.encode('utf-8'),
                                                                    auto_reboot)

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def load_model_from_flash(device_group: DeviceGroup) -> ModelNefDescriptor:
        """
        Load model from device flash (Please update NEF model in flash by Kneron DFUT).

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.

        Returns
        -------
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.ModelNefDescriptor
        """
        model_nef_descriptor = ModelNefDescriptor()
        status = KPWrapper().LIB.kp_load_model_from_flash(device_group.address,
                                                          ctypes.byref(model_nef_descriptor._get_element_buffer()))
        model_nef_descriptor._is_allocate_from_c = True

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

        return model_nef_descriptor

    @staticmethod
    def install_driver_for_windows(product_id: ProductId) -> None:
        """
        Install device driver on Windows

        Parameters
        ----------
        product_id : kp.ProductId
            enum for USB PID(Product ID).

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.ProductId
        """
        status = KPWrapper().LIB.kp_install_driver_for_windows(product_id.value)

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def store_ddr_management_attributes(device_group: DeviceGroup,
                                        ddr_manage_attributes: DdrManageAttributes) -> None:
        """
        (Advance) Store DDR memory management attributes into DeviceGroup.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            A set of devices handle.
        ddr_manage_attributes : kp.DdrManageAttributes
            DDR memory management descriptor of Kneron device.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.DdrManageAttributes

        Notes
        -----
        Must reset-reboot device before setting the DDR memory management attributes.
        Must issue kp.core.store_ddr_management_attributes before kp.core.load_model()/kp.core.load_model_from_file()/kp.core.load_encrypted_models()/kp.core.load_encrypted_models_from_file()/kp.core.load_model_from_flash
        """
        status = KPWrapper().LIB.kp_store_ddr_manage_attr(device_group.address,
                                                          ddr_manage_attributes._get_element_buffer())

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def get_version() -> str:
        """
        Get Kneron PLUS version

        Returns
        -------
        plus_version_string : str
            Kneron PLUS version string.
        """
        return KPWrapper().LIB.kp_get_version().decode('utf-8')
